// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: tianqiang

#pragma once

#include "coin2/exchange/base/api_util/api_signature.h"
#include "coin2/exchange/base/api_util/json_writer.h"

namespace {

std::string GenSignature(
    const std::string& timestamp,
    const std::string& memo,
    const std::string& query_string,
    const std::string& secret_key) {
  std::string before_sign = timestamp + "#" + memo + "#" + query_string;

  return ApiSignatureUtil::GetSignatureByHexHmacSHA256(before_sign, secret_key);
}

}  // namespace

class BitmartWsAuth {
 public:
  explicit BitmartWsAuth(const std::string& key_filepath) : key_(AuthKey::FromFile(key_filepath)) {}

  std::string CreateSignature() const {
    // {"op":"login","args":["<API_KEY>", "<timestamp>", "<sign>"]}

    std::list<std::string> out;
    std::string timestamp = std::to_string(GetCurrentTimestamp() / 1000000);
    std::string sign =
        GenSignature(timestamp, key_.key_name(), "bitmart.WebSocket", key_.secret_key());
    out.emplace_back(key_.access_key());
    out.emplace_back(timestamp);
    out.emplace_back(sign);

    JsonWriter writer;
    writer.PutString("op", "login");
    writer.PutArray("args", out);
    return writer.ToJsonString();
  }

 private:
  AuthKey key_;
};

class BitmartRestAuth : public BaseRestAuth {
 public:
  explicit BitmartRestAuth(
      const std::string& key_filepath,
      const std::string& host,
      const std::string& port,
      const std::string& bind_ip)
      : key_(AuthKey::FromFile(key_filepath)), host_(host), port_(port), bind_ip_(bind_ip) {
    header_.emplace_back(std::make_pair("Content-Type", "application/json"));
    header_.emplace_back(std::make_pair("X-BM-KEY", key_.access_key()));
    header_.emplace_back(std::make_pair("X-BM-SIGN", ""));
    header_.emplace_back(std::make_pair("X-BM-TIMESTAMP", ""));
  }

  explicit BitmartRestAuth(
      const std::string& key_filepath,
      const std::string& host,
      int32_t port,
      const std::string& bind_ip)
      : BitmartRestAuth(key_filepath, host, std::to_string(port), bind_ip) {}

  AsioHttpContext* CreateRequestByGetWithSignature(
      const std::string& url_path,
      UrlParamsBuilder* builder) override {
    std::string query_string = builder->ComputeQueryString();
    std::string post_body = builder->ComputePostBodyAsQueryString();
    CHECK(post_body.empty());

    const auto& header = GenHeader(query_string);

    std::string target = url_path;
    if (!query_string.empty()) {
      target += "?" + query_string;
    }

    return AsioHttpContext::NewGet(host_, port_, bind_ip_, target, header);
  }

  AsioHttpContext* CreateRequestByPostWithSignature(
      const std::string& url_path,
      UrlParamsBuilder* builder) override {
    std::string query_string = builder->ComputeQueryString();
    CHECK(query_string.empty());
    std::string post_body = builder->ComputePostBody();

    const auto& header = GenHeader(post_body);

    return AsioHttpContext::NewPost(host_, port_, bind_ip_, url_path, header, post_body);
  }

  AsioHttpContext* CreateRequestByDeleteWithSignature(
      const std::string& url_path,
      UrlParamsBuilder* builder) override {
    return nullptr;
  }

  AsioHttpContext* CreateRequestByPutWithSignature(
      const std::string& address,
      UrlParamsBuilder* builder) override {
    return nullptr;
  }

  std::vector<std::pair<std::string, std::string>>& GenHeader(const std::string& query_string) {
    const std::string timestamp = std::to_string(GetCurrentTimestamp() / 1000000);
    header_[2].second = GenSignature(timestamp, key_.key_name(), query_string, key_.secret_key());
    header_[3].second = timestamp;

    return header_;
  }

  const std::string& GetAccessKey() { return key_.access_key(); }

 protected:
  AuthKey key_;
  std::string host_;
  std::string port_;
  std::string bind_ip_;
  std::vector<std::pair<std::string, std::string>> header_;
};
